import os
import operator
import asyncio
from dotenv import load_dotenv
from typing import Annotated, Any, TypedDict, List

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END

# --- 1. 环境准备 ---
load_dotenv()

# --- 2. 状态定义 ---
# 我们需要跟踪故事草稿、评审历史以及修改次数
class AgentState(TypedDict):
    topic: str
    draft: str
    critiques: Annotated[list[str], operator.add]
    revision_number: int

# --- 3. 代理节点定义 ---

# a. 通用 LLM
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))

# b. 作者代理 (Author Agent)
author_prompt = PromptTemplate.from_template(
    """你是一位才华横溢的科幻小说作家。你的任务是根据给定的主题创作一个引人入胜的短篇故事。

主题: {topic}
{critiques_str}
你的故事草稿:"""
)

async def author_node(state: AgentState):
    print(f"\n---作者: 第 {state['revision_number']} 次创作---")
    critiques_str = ""
    if state.get("critiques"):
        critiques_str = "请根据以下评论家的意见进行修改:\n- " + "\n- ".join(state["critiques"])
    
    author_chain = author_prompt | llm
    draft = await author_chain.ainvoke({
        "topic": state["topic"],
        "critiques_str": critiques_str,
    })
    
    return {"draft": draft.content}

# c. 评论家代理 (Critic Agent)
critic_prompt = PromptTemplate.from_template(
    """你是一位严谨的文学评论家。你的任务是审查以下故事草稿，并决定它是否可以直接发表。

故事草稿:
{draft}

你的评审:
1.  **评价**: (简要评价故事的优点和缺点)
2.  **决策**: (根据评价，做出 '可以发表' 或 '需要修改' 的决策)
3.  **修改意见**: (如果需要修改，请提供清晰、具体的修改建议)
"""
)

async def critic_node(state: AgentState):
    print("\n---评论家: 审查中...---")
    critic_chain = critic_prompt | llm
    critique_text = await critic_chain.ainvoke({"draft": state["draft"]})
    
    print(f"评论家意见:\n{critique_text.content}")
    
    # 解析评论家的决策
    if "可以发表" in critique_text.content:
        return {
            "critiques": [], # 清空评论，表示满意
        }
    else:
        # 只提取修改意见部分
        try:
            suggestions = critique_text.content.split("修改意见:")[1].strip()
        except IndexError:
            suggestions = "评论家要求修改，但未提供明确意见。"
            
        return {
            "critiques": [suggestions], # 将新意见加入列表
        }

# --- 4. 条件边与图构建 ---

MAX_REVISIONS = 3
def should_continue(state: AgentState) -> str:
    print("---决策: 是否需要继续修改?---")
    
    # 增加修订版本号
    revision_number = state.get("revision_number", 0) + 1
    
    if not state.get("critiques"):
        print(f"决策: 评论家满意，流程结束。")
        return END
    
    if revision_number > MAX_REVISIONS:
        print(f"决策: 已达到最大修改次数 ({MAX_REVISIONS})，流程终止。")
        return END
        
    print(f"决策: 需要第 {revision_number} 次修改。")
    return "author"

# 构建图
workflow = StateGraph(AgentState)
workflow.add_node("author", author_node)
workflow.add_node("critic", critic_node)

workflow.set_entry_point("author")

workflow.add_edge("author", "critic")
workflow.add_conditional_edges(
    "critic",
    should_continue,
    {
        "author": "author",
        END: END,
    },
)

# --- 5. 编译与运行 ---
app = workflow.compile()

async def main():
    initial_state = {
        "topic": "一只生活在赛博朋克城市的猫，它发现了一个能改变现实的秘密。",
        "critiques": [],
        "revision_number": 0,
    }
    
    final_state = await app.ainvoke(initial_state)
    
    print("\n\n---流程结束---")
    print("最终版故事:")
    print(final_state["draft"])

if __name__ == "__main__":
    asyncio.run(main()) 